import cv2
import numpy as np
import pickle


def track_show(image, save_path='hsv_record.pkl'):
    """
    :param image: 输入需要调整hsv值得图像
    :return:
    """
    # 建一个假图像，目的是为了显示一个滚动条的窗口
    tracker = np.zeros((300, 512, 3), np.uint8)
    # 建立一个名为congtrol的窗口
    cv2.namedWindow('control')

    # 创建六个滚动条
    cv2.createTrackbar('Hlow', 'control', 0, 255, nothing)
    cv2.createTrackbar('Hhigh', 'control', 150, 255, nothing)
    cv2.createTrackbar('Slow', 'control', 0, 255, nothing)
    cv2.createTrackbar('Shigh', 'control', 150, 255, nothing)
    cv2.createTrackbar('Vlow', 'control', 0, 255, nothing)
    cv2.createTrackbar('Vhigh', 'control', 150, 255, nothing)
    hsv_list = []
    # 进入循环
    while True:

        # 获取滚动条的动态值
        hl = cv2.getTrackbarPos('Hlow', 'control')
        hh = cv2.getTrackbarPos('Hhigh', 'control')
        sl = cv2.getTrackbarPos('Slow', 'control')
        sh = cv2.getTrackbarPos('Shigh', 'control')
        vl = cv2.getTrackbarPos('Vlow', 'control')
        vh = cv2.getTrackbarPos('Vhigh', 'control')

        # color masking: define color range
        # hsv 下限
        lower = np.array([hl, sl, vl], dtype=np.uint8)
        # hsv 上限
        upper = np.array([hh, sh, vh], dtype=np.uint8)
        # 根据hsv参数生成二值图
        mask = cv2.inRange(image, lower, upper)
        # 显示二值图片
        cv2.imshow("mask", mask)

        # 显示输入的图像
        cv2.imshow('original', image)

        # 显示滚动条窗口
        cv2.imshow('control', tracker)

        # 等待按下q退出
        key = cv2.waitKey(1)
        if key == ord('q') or key == 27:
            with open(save_path, 'wb') as f:
                pickle.dump(hsv_list, f)
                print(f"HSV value write in {save_path}")
            break
        if key == 32:
            hsv_list.append(np.vstack((lower, upper)))
            print(f"hsv value write to list,index:{len(hsv_list) - 1}")


# 滚动条回调函数
def nothing(x):
    pass


def get_hsv_from_file(file_path='hsv_record.pkl'):
    with open(file_path, 'rb') as f:
        hsv_list = pickle.load(f)
        return hsv_list


if __name__ == '__main__':
    image = cv2.imread("images/color_recognition/rect4.jpg")
    track_show(image)
    # ls = get_hsv_from_file('./hsv_record.pkl')
    # l, u = ls[0]
    # mask = cv2.inRange(image, l, u)
    # canny = cv2.Canny(mask, 100, 150)
    # cv2.imshow("canny", canny)
    # cv2.imshow("test", mask)
    # cv2.waitKey(0)